SageMaker Python SDKのローカルモードを利用して、ノートブックインスタンス以外の環境で学習ジョブを回してみる
どうも、大阪DI部の大澤です。
Amazon SageMakerはAWS上でノートブックインスタンスによるデータ探索から機械学習モデルの学習から推論用エンドポイントへのモデルのデプロイなどを行うことができる、機械学習のフルマネージドサービスです。エントリーポイントとなるスクリプトやコンテナイメージを指定することで、組み込みアルゴリズム以外のアルゴリズム(モデル）での学習ジョブやモデルのホスティングを行うことができます。 今回はそんなエントリポイント用スクリプトや学習や推論に使用するコンテナイメージを開発する際に便利な、SageMaker Python SDKの機能であるローカルモードをノートブックインスタンス以外の環境で動かす方法について紹介します。
概要
Irisデータの分類モデルをSageMaker Python SDKのローカルモードを利用して、MacBook Pro上で学習させます。
- モデル: TensorFlowのDNNClassifier
- amazon-sagemaker-examples/iris_dnn_classifier.py at master · awslabs/amazon-sagemaker-examples
- データセット: Iris
- 機械学習や統計分析などのチュートリアルで良く使われるデータセットで、萼片の長さと幅、花びらの長さと幅と花の種類(3種類)のデータが入っています。
やってみる
準備
SageMaker Python SDKのローカルモードではローカルでDockerコンテナを作成し、そのコンテナ内で学習ジョブやモデルのホスティングなどを行います。 そのために必要となるDockerのインストールが必要です。
ローカルモードでローカル環境のGPUを使う場合にはnvidia-dockerのインストールが必要となります。(Macではnividia-docker未対応)
Boto3とSageMaker Python SDKをインストールします。
pip install boto3, sagemaker
学習などの操作で使用するSageMakerのセッションやS3バケット、モデルアーティファクトの保存場所などの設定を行います。 ここで気をつけなければいけないのが、ローカルモードを使用する場合にはローカル用のセッションを使う必要があるということです。通常のSageMakerセッションをEstimatorに渡したり、エンドポイントの削除に使用したりするとエラーが発生します。
import boto3
from sagemaker import get_execution_role
from sagemaker.local import LocalSession
# 今回の操作で使用するプロファイルを指定する(AWS CLIで設定してあるプロファイル)
boto_sess = boto3.Session(profile_name='profile_name', region_name='ap-northeast-1')
# SageMakerの操作で使用するローカル用セッション
local_session = LocalSession(boto_sess)
# 使用するS3のバケット名
bucket = 'バケット名'
# エントリポイントスクリプトのアップロード場所
custom_code_upload_location = 's3://{}/customcode/tensorflow_iris'.format(bucket)
# モデルアーティファクト(学習による生成物)の保存場所
model_artifacts_location = 's3://{}/artifacts'.format(bucket)
# SageMakerでの学習ジョブやモデルのホスティングに使用するIAMロールを指定する
# ローカルモードでは使用されませんが、存在しているIAMロールの指定が必要
role = 'arn:aws:iam::12345:role/sagemaker-execution-role-hogehoge'
学習
ローカルモードで学習ジョブを実行するためにTensorFlow用Estimatorを作成します。 基本的にはSageMakerで通常通り学習ジョブを実行するのと変わらないですが、2点注意する必要があります。
- train_instance_type:
localもしくはlocal_gpuを指定する必要があります。(通常はml.p2.xlarge等のインスタンスタイプを指定) - 今回は
localを指定してCPUでの学習を行いますが、local_gpuを指定することでローカルのGPUを使って学習します。 - sagemaker_session: 準備のところでもお伝えした通り、ローカル用のセッションを指定する必要があります。
- セッションを何も指定しない場合には、内部で自動的にローカル用セッションが作成されるため、問題ありません。
- train_instance_typeでローカルモードを指定しているにも関わらず、sagemaker.session.Session()を渡すとインスタンスタイプが不正だというエラーが出ます。
from sagemaker.tensorflow import TensorFlow
# 学習をハンドルするTensorFlow用Estimator
iris_estimator = TensorFlow(entry_point='iris_dnn_classifier.py',
role=role,
framework_version='1.12.0',
output_path=model_artifacts_location,
code_location=custom_code_upload_location,
train_instance_count=1,
train_instance_type='local', # ローカルモードを指定
training_steps=11,
evaluation_steps=1,
sagemaker_session=local_session) # ローカル用セッションを渡す
# 学習データはSageMakerのサンプル用に公開されているIrisデータセットを使う
region = boto_sess.region_name
train_data_location = 's3://sagemaker-sample-data-{}/tensorflow/iris'.format(region)
# 学習開始
iris_estimator.fit(train_data_location)
SageMakerのプラットフォームで学習させる場合と異なり、ローカルにコンテナを起動させるため、オーバーヘッドがかなり少ないです。ステップ数をかなり少なく設定して動作確認に用いると良さそうです。
※ エントリーポイントとして用いたスクリプト: amazon-sagemaker-examples/iris_dnn_classifier.py at master · awslabs/amazon-sagemaker-examples
モデルのホスティング
先ほど学習させたモデルをローカルモードでホスティングします。
こちらも先ほどの学習と同様にinstance_typeにlocalもしくはlocal_gpuを指定する必要があります。
iris_predictor = iris_estimator.deploy(initial_instance_count=1,
instance_type='local')
こちらも数秒から数分程度でホスティングが完了します。
SageMaker上でモデルをホスティングしている場合と同様にローカルモードでも推論できます。 なので、推論処理についても簡単に動作確認することができます。
iris_predictor.predict([6.4, 3.2, 4.5, 1.5]) #expected label to be 1
エンドポイントの削除もローカルモード用セッションを通じて出来ます。
local_session.delete_endpoint(iris_predictor.endpoint) # ↓でもOK # iris_predictor.delete_endpoint()
さいごに
SageMaker Python SDKのローカルモードを使うことでローカル環境でモデルの学習とホスティングが可能となります。今回紹介した機能以外にもバッチ変換も同様にローカルモードで動かすことが出来ます。SageMakerを使うことでインフラのことを考えることなく、学習やホスティングが可能となりますが、数分のオーバーヘッドが生じます。開発時などで色々試す際にはローカル環境でパッと動かすことができるこの機能は非常に便利だと思いました。
お読みくださり、ありがとうございました-!
